package com.haibara.codesandbox.core;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.command.StatsCmd;
import com.github.dockerjava.api.exception.NotFoundException;
import com.github.dockerjava.api.model.*;
import com.haibara.codesandbox.constant.FileConstant;
import com.haibara.codesandbox.constant.KafkaConstant;
import com.haibara.codesandbox.kafka.processor.Processor;
import com.haibara.codesandbox.model.*;
import com.haibara.codesandbox.model.enums.JudgeResultEnum;
import com.haibara.codesandbox.threadPool.RunCodeTask;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;

import java.io.Closeable;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 代码沙箱模板类
 *
 * @Author: mengyu
 * @Date: 2025/2/12
 */
@Component
@Slf4j
public class DefaultCodeSandbox extends CodeSandboxTemplate {
    /**
     * 文件路径映射：
     * --
     * 系统中路径：
     * tempFilePath = /root/oj/tempFile
     * codePath = /root/oj/code
     */


    @Resource
    private DockerClient dockerClient;
    @Resource
    private Processor processor;
    @Resource
    @Lazy
    private RunCodeTask runCodeTask;

    public static String TIME_EXCEED = "Time Exceed";

    /**
     * 将用户代码存储至对应文件
     *
     * @param code       代码
     * @param questionId 题目id
     * @param userId     用户id
     */
    @Override
    public void saveCodeToFile(String code, Long questionId, Long userId) {
        // 将用户代码存储至对应路径：
        // filePath/{userId}/question{questionId}.sh
        // 如 filePath/1/question1.sh
        // 指定文件路径和名称
        String filePath = FileConstant.codeFilePath + "/" + userId;
        Path scriptPath = Paths.get(filePath, "question" + questionId + ".sh");
        Path parentDir = scriptPath.getParent();
        if (!Files.exists(parentDir)) {
            try {
                Files.createDirectories(parentDir);
            } catch (IOException e) {
                log.error("创建父文件失败 {} ", e.getMessage());
                throw new RuntimeException("创建父文件失败 : " + FileConstant.codeFilePath + "/" + userId);
            }
        }
        try {
            // 将代码写入用户代码文件
            Files.write(scriptPath, Collections.singleton(code), StandardCharsets.UTF_8);
        } catch (IOException e) {
            log.error("将代码写入文件失败 {}", e.getMessage());
            throw new RuntimeException("将代码写入文件失败 : " + filePath + "/" + "question" + questionId + ".sh");
        }
    }

    /**
     * 获取执行指令
     *
     * @param questionCase 测试用例
     * @param questionId   题目id
     * @param userId       用户id
     * @return 指令，其中字符串数组各个元素代表指令中的每个部分，如["/bin/sh","test.sh"]对应指令为 /bin/sh test.sh
     */
    @Override
    public String[] getCmdStrings(QuestionCase questionCase, Long questionId, Long userId) {
        ArrayList<String> cmdStrings = new ArrayList<>();
        // 1. 配置基础指令
        cmdStrings.add("/bin/sh");
        cmdStrings.add(FileConstant.getCodePath(userId, questionId));   // 用户代码文件地址
        // 2. 配置变量输入
        List<String> arguments = JSONUtil.toList(questionCase.getArgumentsInput(), String.class);
        if (arguments != null && !arguments.isEmpty()) {
            cmdStrings.addAll(arguments);
        }
        // 3. 配置文件输入
        List<String> fileNames = JSONUtil.toList(questionCase.getFileInput(), String.class);
        if (fileNames != null && !fileNames.isEmpty()) {
            for (String fileName : fileNames) {
                // 拼接为测试用例文件地址
                cmdStrings.add(FileConstant.getInputFilePath(questionId, questionCase.getId(), fileName));
            }
        }
        // 4. 配置文件输出
        String fileOutput = questionCase.getFileOutput();
        if (!StrUtil.isBlankIfStr(fileOutput)) {
            cmdStrings.add(FileConstant.getUserOutputFilePath(userId, questionId, questionCase.getId()));
            // 清空原本的文件内容
            try {
                Path filePath = Paths.get(FileConstant.getUserOutputFilePath(userId, questionId, questionCase.getId()));
                Files.createDirectories(filePath.getParent());
                Files.write(filePath, new byte[0]);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        return cmdStrings.toArray(new String[0]);
    }

    /**
     * 执行用户代码
     *
     * @param questionCase 参数用例
     * @return 执行结果
     */
    @Override
    public RunCodeResult runCode(
            CreateContainerResponse container, QuestionCase questionCase, Long userId, Long questionId, Long questionSubmitId, Integer timeLimit
    ) throws InterruptedException {
        // 1 验证容器，如果获取容器id失败代表容器未启动
        String containerId = container.getId();

        // 2 设置inputFile为不可变，并给BASH文件消除Windows换行
        ExecCreateCmdResponse clearWindowsSymbol = dockerClient.execCreateCmd(containerId)
                .withCmd("sed", "-i", "s/\\r//", FileConstant.getCodePath(userId, questionId))
                .exec();
        dockerClient.execStartCmd(clearWindowsSymbol.getId())
                .exec(new ResultCallback.Adapter<>())
                .awaitCompletion(5, TimeUnit.SECONDS);

        // 3 配置文件运行命令
        ExecCreateCmdResponse execCreateCmdResponse = dockerClient.execCreateCmd(containerId)
                .withCmd(this.getCmdStrings(questionCase, questionId, userId))
                .withAttachStdout(true)
                .withAttachStderr(true)
                .exec();

        // 4 获取执行输出信息与时间消耗、内存消耗
        long startTime = System.currentTimeMillis();
        //使用同步执行 + CountDownLatch
        CountDownLatch latch = new CountDownLatch(1);
        AtomicReference<ResultMessage> resultRef = new AtomicReference<>();
        AtomicLong timeUsedRef = new AtomicLong();

        ResultCallback.Adapter<Frame> callBack = new ResultCallback.Adapter<>() {
            final StringBuilder stdout = new StringBuilder();
            final StringBuilder stderr = new StringBuilder();

            @Override
            public void onNext(Frame frame) {
                if (frame.getStreamType() == StreamType.STDOUT) {
                    log.info(new String(frame.getPayload()));
                    stdout.append(new String(frame.getPayload()));
                } else {
                    log.info(new String(frame.getPayload()));
                    stderr.append(new String(frame.getPayload()));
                }
            }

            @Override
            public void onComplete() {
                timeUsedRef.set(System.currentTimeMillis() - startTime);
                log.info("容器{}执行完成", containerId);

                if (!stderr.isEmpty()) {
                    resultRef.set(ResultMessage.builder().hasError(true).build());
                } else {
                    resultRef.set(ResultMessage.builder()
                            .runCodeMessage(stdout.toString())
                            .hasError(false)
                            .build());
                }
                latch.countDown(); // 释放锁
                super.onComplete();
            }
        };

        // 内存监控
        final AtomicLong memoryUsed = new AtomicLong(0L);
        StatsCmd statsCmd = dockerClient.statsCmd(containerId);
        ResultCallback<Statistics> statisticsResultCallback = statsCmd.exec(new ResultCallback<Statistics>() {
            @Override
            public void close() {
            }

            @Override
            public void onStart(Closeable closeable) {
            }

            @Override
            public void onNext(Statistics statistics) {
                Long usage = statistics.getMemoryStats().getUsage();
                if (usage != null) {
                    memoryUsed.updateAndGet(current -> Math.max(current, usage));
                }
            }

            @Override
            public void onError(Throwable throwable) {
            }

            @Override
            public void onComplete() {
            }
        });
        statsCmd.exec(statisticsResultCallback);

        // 执行并等待完成
        dockerClient.execStartCmd(execCreateCmdResponse.getId()).exec(callBack);

        // 等待执行完成，设置超时时间
        boolean completed = latch.await(timeLimit / 1000 + 5, TimeUnit.SECONDS);
        if (!completed) {
            // 处理超时情况
            log.warn("代码执行超时: userId:{}, questionId:{}", userId, questionId);
            return RunCodeResult.builder()
                    .memoryUsed(0L)
                    .timeUsed(0L)
                    .resultMessage(
                            ResultMessage.builder()
                                    .hasError(true)
                                    .runCodeMessage(TIME_EXCEED)
                                    .build()
                    ).build();
        }

        log.info("运行完成: userId:{}, questionId:{}", userId, questionId);
        return RunCodeResult.builder()
                .resultMessage(resultRef.get())
                .timeUsed(timeUsedRef.get())
                .memoryUsed(memoryUsed.get() / 1000)
                .build();
    }

    private HostConfig getHostConfig(Long userId, Long questionId) {
        // 2.1.1 配置数据卷
        HostConfig hostConfig = HostConfig.newHostConfig();
        hostConfig.withBinds(
                // 用户代码数据卷
                Bind.parse(FileConstant.codeFilePath + "/" + userId + ":" + FileConstant.codeFilePath + "/" + userId),
                // 测试文件数据卷，设定为只读
                Bind.parse(FileConstant.inputFilePath + "/" + questionId + ":" + FileConstant.inputFilePath + "/" + questionId + ":ro"),
                // 文件输出数据卷
                Bind.parse(FileConstant.userOutputFilePath + "/" + userId + "/" + questionId + ":" + FileConstant.userOutputFilePath + "/" + userId + "/" + questionId)
        );
        // 2.1.2 配置容器限制
        hostConfig
                .withCpuShares(512)
                // 禁止特权命令
                .withPrivileged(false)
                .withCapDrop(Capability.valueOf("ALL"))
                // 设置容器内部文件为只读
                .withReadonlyRootfs(true)
                // 禁止容器联网
                .withNetworkMode("none")
                // 设置最大内存
                .withMemory(262144000L);
        return hostConfig;
    }


    /**
     * 执行用户代码并收集结果
     *
     * @param executeCodeRequest 代码请求
     */
    @Override
    public void executeCode(ExecuteCodeRequest executeCodeRequest) throws InterruptedException, ExecutionException {
        Long questionSubmitId = executeCodeRequest.getQuestionSubmitId();
        List<QuestionCase> questionCases = executeCodeRequest.getQuestionCases();
        String code = executeCodeRequest.getCode();
        Long questionId = executeCodeRequest.getQuestionId();
        Long userId = executeCodeRequest.getUserId();
        Integer timeLimit = executeCodeRequest.getTimeLimit();

        // 1. 将用户代码存储至文件
        this.saveCodeToFile(code, questionId, userId);
        // 2. 执行代码
        // 2.1 启动容器
        CreateContainerResponse container = dockerClient.createContainerCmd("alpine:3.12.0")
                .withTty(true)
                .withAttachStdin(true)
                .withAttachStderr(true)
                .withAttachStdout(true)
                // 配置数据卷
                .withHostConfig(getHostConfig(userId, questionId))
                .exec();
        dockerClient.startContainerCmd(container.getId()).exec();
        log.info("容器id:{}", (container.getId()));
        // 2.2 执行代码
        ArrayList<CompletableFuture<RunCodeResult>> runCodeCompletableFutureResult = new ArrayList<>();
        for (QuestionCase questionCase : questionCases) {
            log.info("执行测试用例 : {}", questionCase.getId());
            runCodeCompletableFutureResult.add(
                    runCodeTask.createTask(container, questionCase, userId, questionId, timeLimit, questionSubmitId)
            );
        }
        // 3. 构建ExecuteCodeRespond
        CompletableFuture.allOf(
                        runCodeCompletableFutureResult.toArray(new CompletableFuture[0])
                ).thenAccept(
                        none -> {
                            // 3.1 获取执行结果
                            List<RunCodeResult> runCodeResultList = runCodeCompletableFutureResult.stream()
                                    .map(CompletableFuture::join).toList();
                            // 3.2 判断执行是否失败
                            for (RunCodeResult runCodeResult : runCodeResultList) {
                                if (runCodeResult.getResultMessage().getHasError()) {
                                    //有错误
                                    String errorMsg = runCodeResult.getResultMessage().getRunCodeMessage();
                                    JudgeResultEnum resultType = TIME_EXCEED.equals(errorMsg)
                                            ? JudgeResultEnum.TIME_LIMIT_EXCEEDED
                                            : JudgeResultEnum.RUNTIME_ERROR;
                                    // 构建错误响应
                                    ExecuteCodeRespond errorResponse = ExecuteCodeRespond.builder()
                                            .outputList(null)
                                            .message(resultType.getText())
                                            .resultStatus(ExecuteCodeRespond.ERROR)
                                            .memoryUsed(0L)
                                            .timeUsed(0L)
                                            .questionSubmitId(questionSubmitId)
                                            .build();
                                    log.info("发送{}的判题请求: {}",
                                            resultType == JudgeResultEnum.TIME_LIMIT_EXCEEDED ? "超时" : "运行错误",
                                            errorResponse);

                                    processor.sendMsg(KafkaConstant.JUDGE_TOPIC, JSONUtil.toJsonStr(errorResponse));
                                    return;
                                }
                            }
                            // 3.3 构建参数输出结果
                            List<String> outputList = new ArrayList<>();
                            long timeUsed = 0L;
                            long memoryUsed = 0L;
                            for (RunCodeResult codeResult : runCodeResultList) {
                                // 消除可能出现的windows换行符
                                outputList.add(codeResult.getResultMessage().getRunCodeMessage().replaceAll("\\r", ""));
                                // 时间消耗和内存消耗选取最大值
                                timeUsed = Math.max(timeUsed, codeResult.getTimeUsed());
                                memoryUsed = Math.max(memoryUsed, codeResult.getMemoryUsed());
                            }
                            // 3.4 构建文件输出结果
                            List<String> fileOutputList = questionCases.stream()
                                    .map(c ->
                                            StrUtil.isBlankIfStr(c.getFileOutput())
                                                    ? ""
                                                    : "case" + c.getId() + ".txt"
                                    )
                                    .toList();
                            // 3.5 发送判题请求
                            ExecuteCodeRespond successECR = ExecuteCodeRespond.builder()
                                    .outputList(outputList)
                                    .fileOutputList(fileOutputList)
                                    .resultStatus(ExecuteCodeRespond.SUCCESS)
                                    .memoryUsed(memoryUsed)
                                    .timeUsed(timeUsed)
                                    .questionSubmitId(questionSubmitId)
                                    .build();
                            processor.sendMsg(KafkaConstant.JUDGE_TOPIC, JSONUtil.toJsonStr(successECR));
                            log.info("发送正常的判题请求:{}", successECR);
                        }
                ).whenComplete((r, ex) -> {
                    // 删除容器
                    log.info("同步删除容器开始{}", container.getId());
                    try {
                        dockerClient.removeContainerCmd(container.getId())
                                .withForce(true)
                                .exec();
                        log.info("同步删除容器结束{}", container.getId());
                    } catch (NotFoundException e) {
                        log.info("同步删除容器失败{}", e.getMessage());
                    }
                })
                .exceptionally(ex -> {
                    // 若发生异常则发送System Error信息
                    log.error("发生异常 :", ex);
                    ExecuteCodeRespond errorECR = ExecuteCodeRespond.builder()
                            .outputList(null)
                            .message(JudgeResultEnum.SYSTEM_ERROR.getText())
                            .resultStatus(ExecuteCodeRespond.SYSTEM_ERROR)
                            .memoryUsed(0L)
                            .timeUsed(0L)
                            .questionSubmitId(questionSubmitId)
                            .build();
                    log.error("发送异常的判题请求:{}", errorECR);
                    processor.sendMsg(KafkaConstant.JUDGE_TOPIC, JSONUtil.toJsonStr(errorECR));
                    return null;
                });
    }
}
